import random
from sequenceAnalysisFunctions import *

######################################################################################
#simple supporting functions
######################################################################################

def get_sequence_of_AAs(sequenceAsList):
    codonAAdict = get_codon_to_AA_dict()
    aaSequenceList = []
    for codon in sequenceAsList[:-1]:
        aaSequenceList.append(codonAAdict[codon])

    return aaSequenceList

def calculate_codon_probabilities(codonDictSums):
    codonDictRanges = get_codon_dict()
    codonDictProbabilities = get_codon_dict()
    aaDict = get_AA_dict()
    for aminoAcid in aaDict.keys():
        totalCodons = 0
        for codon in aaDict[aminoAcid]:
            totalCodons += codonDictSums[codon]
        for codon in aaDict[aminoAcid]:
            try:
                codonDictProbabilities[codon] = codonDictSums[codon] / float(totalCodons)
            except ZeroDivisionError:
                codonDictProbabilities[codon] = 0.
        current = 0.
        for codon in aaDict[aminoAcid]:
            probability = codonDictProbabilities[codon]
            codonDictRanges[codon] = (current, current + probability)
            current = current + probability
    return codonDictRanges

def randomization_procedure(aminoAcidSequence, codonDictSums):
    aaDict = get_AA_dict()
    newString = ''
    for aminoAcid in aminoAcidSequence:
        codonDictRanges = calculate_codon_probabilities(codonDictSums)
        randomNumber = random.random()
        codonList = aaDict[aminoAcid]
        for codon in codonList:
            if randomNumber >= codonDictRanges[codon][0] and randomNumber < codonDictRanges[codon][1]:
                newString = newString + codon
                codonDictSums[codon] -= 1
    if len(newString) != 3*len(aminoAcidSequence):
        print 'ERROR in randomization procedure'
    return newString, codonDictSums

######################################################################################
#main functions to call
######################################################################################

######################################################################################
#this function should randomize a string and preserve codon usage WITHIN that string
def create_random_sequence_individual(sequenceAsString):
    codonDictSums = get_codon_dict()
    sequenceAsList = split_string_into_list_of_triplets(sequenceAsString)
    aminoAcidSequence = get_sequence_of_AAs(sequenceAsList)
    for x in sequenceAsList:
        codonDictSums[x] += 1
    newSequenceAsString, codonDictSums = randomization_procedure(aminoAcidSequence, codonDictSums)
    return newSequenceAsString
######################################################################################
#this is a simple meta function to make multiple calls of the above function
#an entire gene set will be randomized while preserving codon usage WITHIN each gene
def create_random_sequence_individual_set(listOfStrings):
    newListOfStrings = []
    for gene in listOfStrings:
        newListOfStrings.append(create_random_sequence_individual(gene))
    return newListOfStrings



######################################################################################
#this function NEEDS WORK/TO BE TESTED
#I want it to randomize sequences by preserving codon usage of the greater set, but not
#necessarily within each gene
def create_random_sequence_gene_set(listOfStrings):
    codonDictSums = get_codon_dict()
    listOfSequenceLists = [split_string_into_list_of_triplets(gene) for gene in listOfStrings]
    listOfAminoAcidLists = [get_sequence_of_AAs(gene) for gene in listOfSequenceLists]
    for gene in listOfSequenceLists:
        for codon in gene:
            codonDictSums[codon] += 1
    newListOfStrings = []
    for geneAAs in listOfAminoAcidLists:
        newSequenceAsString, codonDictSums = randomization_procedure(geneAAs, codonDictSums)
        newListOfStrings.append(newSequenceAsString)
    return newListOfStrings



######################################################################################
#tests
######################################################################################

